# coding: utf-8

from datetime import datetime, timedelta

from flask_jwt_extended import get_jwt, create_access_token, get_jwt_identity, set_access_cookies

from exts import jwt, rdb
from server.utils.restful import rest_error


# 防止已注销登录的token继续使用
@jwt.token_in_blocklist_loader
def check_if_token_is_revoked(jwt_header, jwt_payload):
    jti = jwt_payload["jti"]
    token_in_redis = rdb.get(jti)
    return token_in_redis is not None


@jwt.expired_token_loader
def my_expired_token_callback(jwt_header, jwt_payload):
    return rest_error("token已过期，请重新登录"), 401


def hook(app):
    @app.after_request
    def after_request(response):
        try:
            exp_timestamp = get_jwt()["exp"]
            now = datetime.now()
            target_timestamp = datetime.timestamp(now + timedelta(minutes=30))
            if target_timestamp > exp_timestamp:
                access_token = create_access_token(identity=get_jwt_identity())
                set_access_cookies(response, access_token)
            return response
        except (RuntimeError, KeyError):
            return response
